package com.kinggm.aop;

import com.kinggm.ano.RedisApiLimit;
import com.kinggm.constants.LimitTypeEnum;
import com.kinggm.utils.IpUtil;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.Collections;

@Aspect
@Component
public class RedisApiLimitAop {


    @Resource
    private HttpServletResponse response;
    @Resource
    private HttpServletRequest request;

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    @Resource
    private DefaultRedisScript<Long> redisLuaScript;


    @Pointcut(value = "@annotation(com.kinggm.ano.RedisApiLimit)")
    public void rateLimit() {
    }


    @Before("rateLimit()")
    public void interceptor(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        RedisApiLimit rateLimit = method.getAnnotation(RedisApiLimit.class);
        if (rateLimit == null) {
            return;
        }

        String key = getKeyByLimitType(rateLimit, signature);

        //调用lua脚本，获取返回结果，这里即为请求的次数
        Long number = redisTemplate.execute(redisLuaScript, Collections.singletonList(key), rateLimit.count(), rateLimit.period());
        if (number != null && number.intValue() != 0 && number.intValue() <= rateLimit.count()) {
            System.out.println("限流时间段内访问了第：" + number + " 次");
            return;
        }

        System.out.println("访问频率过快，被限流了");
        dealResponse(response, "访问频率过快，被限流了");
    }


    /**
     * redis key 有三种
     * 1. 自定义Key: prefix + ":" + key (key 不能为空)
     * 2. 接口key：prefix + ":" + 接口全类名
     * 3. Ip key：prefix + ":" + ip + "-" + 接口全类名
     */
    private String getKeyByLimitType(RedisApiLimit redisApiLimit, MethodSignature signature) {
        String key = "REDIS_API_LIMIT_";
        LimitTypeEnum limitTypeEnum = redisApiLimit.limitType();
        if (LimitTypeEnum.CUSTOMER == limitTypeEnum) {
            // 自定义
            String tempKey = redisApiLimit.key();
            if (!StringUtils.hasText(tempKey)) {
                throw new RuntimeException("自定义类型下 key 不能为空!");
            }
            return key + tempKey;
        }
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        String classFullName = targetClass.getName() + "-" + method.getName();
        if (LimitTypeEnum.INTERFACE == limitTypeEnum) {
            return key + classFullName;
        }
        // IP
        String ipAddress = IpUtil.getIpAddr(request);
        return key + ipAddress + "-" + classFullName;
    }

    private void dealResponse(HttpServletResponse resp, String result) {
        resp.setContentType("application/json;charset=UTF-8");
        try (ServletOutputStream outputStream = resp.getOutputStream()) {
            outputStream.write(result.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


}
